from Node import Node as dotN
from Node import Point
import numpy as np
import random
import Plotmaze
import pickle


def nextPoint(point, move):  # 根据动作生成下一个点的坐标，三维空间的26个方向（此处供外部使用，可不用实例化）
    temppoint = Point(point.x, point.y, point.z)
    if move == 0:
        temppoint.z = point.z + 1
    elif move == 1:
        temppoint.z = point.z + 1
        temppoint.y = point.y + 1
    elif move == 2:
        temppoint.z = point.z + 1
        temppoint.y = point.y - 1
    elif move == 3:
        temppoint.z = point.z + 1
        temppoint.x = point.x + 1
    elif move == 4:
        temppoint.z = point.z + 1
        temppoint.x = point.x - 1
    elif move == 5:
        temppoint.z = point.z + 1
        temppoint.y = point.y + 1
        temppoint.x = point.x + 1
    elif move == 6:
        temppoint.z = point.z + 1
        temppoint.y = point.y + 1
        temppoint.x = point.x - 1
    elif move == 7:
        temppoint.z = point.z + 1
        temppoint.y = point.y - 1
        temppoint.x = point.x + 1
    elif move == 8:
        temppoint.z = point.z + 1
        temppoint.y = point.y - 1
        temppoint.x = point.x - 1
    elif move == 9:
        temppoint.y = point.y + 1
    elif move == 10:
        temppoint.y = point.y - 1
    elif move == 11:
        temppoint.x = point.x + 1
    elif move == 12:
        temppoint.x = point.x - 1
    elif move == 13:
        temppoint.x = point.x + 1
        temppoint.y = point.y + 1
    elif move == 14:
        temppoint.x = point.x + 1
        temppoint.y = point.y - 1
    elif move == 15:
        temppoint.x = point.x - 1
        temppoint.y = point.y + 1
    elif move == 16:
        temppoint.x = point.x - 1
        temppoint.y = point.y - 1
    elif move == 17:  # 正下
        temppoint.z = point.z - 1
    elif move == 18:  # 斜下
        temppoint.z = point.z - 1
        temppoint.y = point.y + 1
    elif move == 19:
        temppoint.z = point.z - 1
        temppoint.y = point.y - 1
    elif move == 20:
        temppoint.z = point.z - 1
        temppoint.x = point.x + 1
    elif move == 21:
        temppoint.z = point.z - 1
        temppoint.x = point.x - 1
    elif move == 22:
        temppoint.z = point.z - 1
        temppoint.y = point.y + 1
        temppoint.x = point.x + 1
    elif move == 23:
        temppoint.z = point.z - 1
        temppoint.y = point.y + 1
        temppoint.x = point.x - 1
    elif move == 24:
        temppoint.z = point.z - 1
        temppoint.y = point.y - 1
        temppoint.x = point.x + 1
    elif move == 25:
        temppoint.z = point.z - 1
        temppoint.y = point.y - 1
        temppoint.x = point.x - 1
    else:
        print('wrong act type')
        return

    nextpoint = temppoint
    return nextpoint


class Map:
    def __init__(self, length, width, height, startPoint, endPoint, obstacle, greedy, rate):
        self.length = length
        self.width = width
        self.height = height  # 地图信息
        self.startNode = startPoint
        self.endNode = endPoint  # 起点终点
        self.greedy = greedy  # 不发生贪婪的概率
        self.learningrate = rate  # 学习率
        self.countrate = 0.3  # 折扣量
        self.obstacle = obstacle  # 障碍物信息
        self.Nodedict = {}  # 所有结点的字典

    def initMap(self):  # 对可行空间的每个坐标均生成结点，生成key为坐标，value为结点的字典
        for i in range(self.length):
            for j in range(self.width):
                for k in range(self.height):
                    temppoint = (i, j, k)
                    makepoint = Point(i, j, k)
                    tempnode = dotN(makepoint, self.length, self.width, self.height, self.obstacle)
                    # 生成结点
                    tempnode.isOut()
                    tempnode.isObstacle()  # 判断结点类型
                    self.Nodedict[temppoint] = tempnode  # 通过坐标建立字典索引结点
        return self.Nodedict

    def resetR(self):  # 继续初始化地图信息，更改部分结点的r、Q值
        for point in self.Nodedict:
            position = self.Nodedict[point].point
            count = 0
            for i in range(26):
                nextposition = self.nextPoint(position, i)
                if 0 <= nextposition.x < self.length and 0 <= nextposition.y < self.width and 0 <= nextposition.z < self.height:
                    nextpoint = (nextposition.x, nextposition.y, nextposition.z)
                    nextnode = self.Nodedict[nextpoint]
                    if nextnode.r == -100:  # 下一步是地图内的点，且该点奖励为-100（即障碍物）
                        count = count + 1  # 不可行数量count增加
                else:
                    count = count + 1  # 下一步不在地图内，不可行数量count增加
            self.Nodedict[point].r = 0 - count  # 该点的奖励为周围障碍物或不可行点数量的负值，减少路途
            if self.Nodedict[point].isObstacle():  # 该点本身为障碍物点时，更新Q值与r值
                self.Nodedict[point].r = -100
                for i in range(26):
                    self.Nodedict[point].Q = np.zeros(26) - 500
            if self.Nodedict[point].isEnd(self.endNode):  # 该点为终点时，更新Q值与r值
                self.Nodedict[point].r = 1000
                for i in range(26):
                    self.Nodedict[point].Q = np.zeros(26) + 1000

    def nextPoint(self, point, move):  # 根据动作生成下一个点的坐标，三维空间的26个方向
        temppoint = Point(point.x, point.y, point.z)
        if move == 0:  # 正上
            temppoint.z = point.z + 1
        elif move == 1:  # 斜上
            temppoint.z = point.z + 1
            temppoint.y = point.y + 1
        elif move == 2:
            temppoint.z = point.z + 1
            temppoint.y = point.y - 1
        elif move == 3:
            temppoint.z = point.z + 1
            temppoint.x = point.x + 1
        elif move == 4:
            temppoint.z = point.z + 1
            temppoint.x = point.x - 1
        elif move == 5:
            temppoint.z = point.z + 1
            temppoint.y = point.y + 1
            temppoint.x = point.x + 1
        elif move == 6:
            temppoint.z = point.z + 1
            temppoint.y = point.y + 1
            temppoint.x = point.x - 1
        elif move == 7:
            temppoint.z = point.z + 1
            temppoint.y = point.y - 1
            temppoint.x = point.x + 1
        elif move == 8:
            temppoint.z = point.z + 1
            temppoint.y = point.y - 1
            temppoint.x = point.x - 1
        elif move == 9:  # 平动
            temppoint.y = point.y + 1
        elif move == 10:
            temppoint.y = point.y - 1
        elif move == 11:
            temppoint.x = point.x + 1
        elif move == 12:
            temppoint.x = point.x - 1
        elif move == 13:
            temppoint.x = point.x + 1
            temppoint.y = point.y + 1
        elif move == 14:
            temppoint.x = point.x + 1
            temppoint.y = point.y - 1
        elif move == 15:
            temppoint.x = point.x - 1
            temppoint.y = point.y + 1
        elif move == 16:
            temppoint.x = point.x - 1
            temppoint.y = point.y - 1
        elif move == 17:  # 正下
            temppoint.z = point.z - 1
        elif move == 18:  # 斜下
            temppoint.z = point.z - 1
            temppoint.y = point.y + 1
        elif move == 19:
            temppoint.z = point.z - 1
            temppoint.y = point.y - 1
        elif move == 20:
            temppoint.z = point.z - 1
            temppoint.x = point.x + 1
        elif move == 21:
            temppoint.z = point.z - 1
            temppoint.x = point.x - 1
        elif move == 22:
            temppoint.z = point.z - 1
            temppoint.y = point.y + 1
            temppoint.x = point.x + 1
        elif move == 23:
            temppoint.z = point.z - 1
            temppoint.y = point.y + 1
            temppoint.x = point.x - 1
        elif move == 24:
            temppoint.z = point.z - 1
            temppoint.y = point.y - 1
            temppoint.x = point.x + 1
        elif move == 25:
            temppoint.z = point.z - 1
            temppoint.y = point.y - 1
            temppoint.x = point.x - 1
        else:
            print('wrong act type')
            return

        nextpoint = temppoint
        return nextpoint

    def run(self, maxstepnum, learningtimes):
        pickle.dump(self.obstacle, open('../pkl/obstacle.pkl', 'wb'))  # 存储障碍物信息
        stepNum = 0  # 计量总步数
        epoch = 0  # 训练中跑到一次终点为一个epoch
        steplist = []
        if stepNum > maxstepnum:
            print("没有发现通路")
            return

        currpoint = self.startNode  # 当前结点为初始点时
        currnode = self.Nodedict[currpoint]
        steplist.append(currpoint)
        currnode.fathernode = currnode

        # print(currnode.Q)

        while stepNum <= learningtimes:
            writeflag = 0  # 是否可达
            if random.uniform(0, 1) >= self.greedy:  # 贪婪情况
                maxQ = np.max(currnode.Q)  # 当前结点所有动作的最大Q值
                # actlist = np.where(currnode.Q == maxQ)[0]
                # lenlist = len(actlist)
                # maxact = actlist[random.randint(0, lenlist - 1)]
                maxact = np.where(currnode.Q == maxQ)[0][0]  # 选取第一个最大的Q值动作作为执行动作
                # print("currpoint",currnode.point.x,currnode.point.y,currnode.point.z)
                nextpointxyz = self.nextPoint(currnode.point, maxact)  # 得到下一个状态的位置
                # print("nextpointxyz",nolintnextline.z)
                if 0 <= nextpointxyz.x < self.length and 0 <= nextpointxyz.y < self.width and 0 <= nextpointxyz.z < self.height:
                    # 在地图空间内时
                    nextpoint = (nextpointxyz.x, nextpointxyz.y, nextpointxyz.z)
                    nextnode = self.Nodedict[nextpoint]
                    outflag = nextnode.isOut()
                    obstacleflag = nextnode.isObstacle()
                    totalflag = outflag + obstacleflag  # 两者只要有一个，便为不可行
                    # print("yes?")
                else:
                    totalflag = 1  # 不在地图内则不可行
                    # currnode.Q[maxact] = currnode.Q[maxact] + self.learningrate * (
                    #        -100 + self.countrate * (-100) - currnode.Q[maxact])
                    currnode.Q[maxact] = currnode.Q[maxact] + self.learningrate * (
                        -100 + self.countrate * (-100) - currnode.Q[maxact]
                    )

                # print(totalflag)
                if totalflag == 0:  # 如果为可行，更新Q值表
                    currnode.Q[maxact] = currnode.Q[maxact] + self.learningrate * (
                            nextnode.r + self.countrate * np.max(nextnode.Q) - currnode.Q[maxact])
                    nextnode.fathernode = currnode
                    currnode = nextnode

            else:  # 非贪婪情况，随机选取动作，并进行上述过程
                move = random.randint(0, 25)
                nextpointxyz = self.nextPoint(currnode.point, move)
                if 0 <= nextpointxyz.x < self.length and 0 <= nextpointxyz.y < self.width and 0 <= nextpointxyz.z < self.height:
                    nextpoint = (nextpointxyz.x, nextpointxyz.y, nextpointxyz.z)
                    nextnode = self.Nodedict[nextpoint]
                    outflag = nextnode.isOut()
                    obstacleflag = nextnode.isObstacle()
                    totalflag = outflag + obstacleflag
                else:
                    totalflag = 1
                    # currnode.Q[move] = currnode.Q[move] + self.learningrate * (
                    #        -100 + self.countrate * (-100) - currnode.Q[move])
                    #currnode.Q[move] = currnode.Q[move] - 500
                    currnode.Q[move] = currnode.Q[move] + self.learningrate * (
                                -100 + self.countrate * (-100) - currnode.Q[move])

                if totalflag == 0:
                    currnode.Q[move] = currnode.Q[move] + self.learningrate * (
                            nextnode.r + self.countrate * np.max(nextnode.Q) - currnode.Q[move])
                    nextnode.fathernode = currnode
                    currnode = nextnode

            steplist.append((currnode.point.x, currnode.point.y, currnode.point.z))  # 将下一个节点加入路径列表中
            # Plotmaze.plotpoint(self.obstacle, self.startNode, self.endNode, steplist)
            stepNum = stepNum + 1  # 增加总步数

            if currnode.isEnd(self.endNode):  # 如果当前结点为终点，则可写并且epoch增加，重置curr指针，清空训练路径点
                epoch = epoch + 1
                writeflag = 1
                currnode = self.Nodedict[self.startNode]
                steplist.clear()
                steplist.append(self.startNode)
            elif stepNum % 1000 == 0:  # 每训练1000步仍没有找到终点的，也视为结束并且epoch增加
                epoch = epoch + 1
                writeflag = 1
                currnode = self.Nodedict[self.startNode]
                steplist.clear()
                steplist.append(self.startNode)

            if epoch % 10000 == 0 and writeflag == 1:  # 每10000个epoch打印一次信息
                print("epoch:", epoch)
                print("newpoint:", (currnode.point.x, currnode.point.y, currnode.point.z))
                print("nextQ:", currnode.Q)

            if epoch % 100000 == 0 and writeflag == 1:  # 每100000个epoch存储一次
                epochstr = str(epoch)
                part1 = '../pkl/Nodedict'
                part2 = '.pkl'
                pathname = part1 + epochstr + part2
                pickle.dump(self.Nodedict, open(pathname, 'wb'), protocol=4)

            if epoch >= 2700001:  # epoch足够大时结束训练
                break

            # if stepNum == 2:
            #    break
            # Plotmaze.plotpoint(obstacle, startPoint, endPoint,steplist)
